import json
import os
import sys
import shutil
import itertools # 导入itertools用于参数组合
from datetime import datetime, timedelta
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed

import pandas as pd
import backtrader as bt
import baostock as bs
from tqdm import tqdm
from pandas.tseries.offsets import BDay

# --- 从新模块中导入组件 ---
from . import data_provider
from . import signal_generator
from .utils import save_report
from .data_provider import get_market_data, prepare_stock_data, filter_sub_new_stocks
from .signal_generator import generate_buy_signals_with_pandas
from .strategy import StockSelectionStrategy, ProgressObserver
# -------------------------

def get_config():
    """
    加载并返回主配置文件 config.json 的内容。
    会从 main.py 的上级目录寻找 config.json
    """
    # 获取当前文件 (main.py) 的目录
    dir_path = os.path.dirname(os.path.realpath(__file__))
    # 构建到上级目录中 config.json 的路径
    config_path = os.path.abspath(os.path.join(dir_path, '..', 'config.json'))
    
    if not os.path.exists(config_path):
        print(f"CRITICAL: config.json not found at expected path: {config_path}")
        raise FileNotFoundError(f"配置文件未找到，路径: {config_path}")

    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    return config

def create_empty_report(momentum_threshold, max_pullback_pct, trailing_stop_pct, trade_size, initial_cash, note):
    """
    创建一个空的回测报告，用于跳过回测的情况
    
    :param momentum_threshold: 动量阈值
    :param max_pullback_pct: 最大回调百分比
    :param trailing_stop_pct: 跟踪止损百分比
    :param trade_size: 交易股数
    :param initial_cash: 初始资金
    :param note: 备注信息
    :return: 包含空回测结果的字典
    """
    return {
        'momentum_threshold': momentum_threshold,
        'max_pullback_pct': max_pullback_pct,
        'trailing_stop_pct': trailing_stop_pct,
        'trade_size': trade_size,
        'Total Return (%)': 0,
        'Sharpe Ratio': 0,
        'Win Rate (%)': 0,
        'Profit/Loss Ratio': 0,
        'Total Trades': 0,
        'Final Portfolio Value': initial_cash,
        'Note': note
    }

def init_baostock_login():
    """
    为每个子进程初始化BaoStock登录。
    """
    print(f"为进程 {os.getpid()} 初始化BaoStock登录...")
    bs.login()
    
def run_backtest(config, stock_codes, market_df, **params):
    """
    执行单次回测的核心逻辑。
    
    :param config: 全局配置字典
    :param stock_codes: 用于生成信号的股票代码列表
    :param market_df: 大盘指数数据
    :param params: 包含所有策略参数的字典
    :return: 包含回测结果的字典
    """
    # --- 从params字典中提取参数，并提供默认值 ---
    default_params = config.get('default_strategy_params', {})
    momentum_threshold = params.get('momentum_threshold', default_params.get('momentum_threshold'))
    max_pullback_pct = params.get('max_pullback_pct', default_params.get('max_pullback_pct'))
    trailing_stop_pct = params.get('trailing_stop_pct', default_params.get('trailing_stop_pct'))
    trade_size = params.get('trade_size', default_params.get('trade_size'))
    ma_divergence_20_max = params.get('ma_divergence_20_max', default_params.get('ma_divergence_20_max'))
    ma_divergence_60_max = params.get('ma_divergence_60_max', default_params.get('ma_divergence_60_max'))
    rsi_overbought_threshold = params.get('rsi_overbought_threshold', default_params.get('rsi_overbought_threshold'))
    precise_pullback_max_dist_pct = params.get('precise_pullback_max_dist_pct', default_params.get('precise_pullback_max_dist_pct'))
    market_rsi_max = params.get('market_rsi_max', default_params.get('market_rsi_max'))

    bt_settings = config['backtest_settings']
    output_settings = config['output_settings']
    data_settings = config['data_settings']
    CACHE_DIR = 'feather_cache'

    start_date_obj = pd.to_datetime(bt_settings['start_date'])
    end_date_obj = pd.to_datetime(bt_settings['end_date'])
    PREHEAT_TRADING_DAYS = 80
    start_date_for_data_obj = start_date_obj - BDay(PREHEAT_TRADING_DAYS)

    # --- 阶段五：使用Pandas生成所有买入信号 ---
    print("\n--- 阶段一：使用Pandas生成买入信号 ---")
    print(f"当前测试参数: {params}")
    
    buy_signals, hot_stocks, hot_stocks_dfs = generate_buy_signals_with_pandas(
        all_stock_codes=stock_codes,
        market_df=market_df.copy(),
        cache_dir=CACHE_DIR,
        start_date_str=bt_settings['start_date'],
        end_date_str=bt_settings['end_date'],
        momentum_threshold=momentum_threshold,
        max_pullback_pct=max_pullback_pct,
        max_price=bt_settings.get('max_price', 9999),
        required_start_date_obj=start_date_for_data_obj.date(),
        ma_divergence_20_max=ma_divergence_20_max,
        ma_divergence_60_max=ma_divergence_60_max,
        rsi_overbought_threshold=rsi_overbought_threshold,
        precise_pullback_max_dist_pct=precise_pullback_max_dist_pct,
        market_rsi_max=market_rsi_max
    )
    
    if not hot_stocks:
        print("\n--- 信号生成阶段未发现任何股票，跳过此次回测。 ---")
        return create_empty_report(
            momentum_threshold, 
            max_pullback_pct, 
            trailing_stop_pct, 
            trade_size, 
            bt_settings['initial_cash'],
            'No signals generated'
        ), 0

    print(f"\n--- 海选完成：共发现 {len(hot_stocks)} 只潜力股，合计产生 {sum(len(v) for v in buy_signals.values())} 个买入信号。 ---")

    # --- 阶段六：使用Backtrader进行精准回测 ---
    print(f"\n--- 阶段二：对海选出的 {len(hot_stocks)} 只股票进行精准回测 ---")
    cerebro = bt.Cerebro(stdstats=False)
    # 设置broker参数，使交易在当天收盘价执行(coc=True)
    cerebro.broker.setcash(bt_settings['initial_cash'])
    cerebro.broker.setcommission(commission=bt_settings['commission_rate'])
    cerebro.broker.set_coc(True)  # 设置为收盘价成交(Close on Close)
    
    strategy_params = {
        'buy_signals': buy_signals,
        'backtest_start_date': bt_settings['start_date'],
        'trailing_stop_pct': trailing_stop_pct,
        'trade_size': trade_size
    }
    # --- 新增：从默认参数中获取长上影线卖出阈值 ---
    strategy_params['sell_on_long_upper_shadow_pct'] = config.get('default_strategy_params', {}).get('sell_on_long_upper_shadow_pct', 0.03)

    cerebro.addstrategy(StockSelectionStrategy, **strategy_params)
    cerebro.addobserver(ProgressObserver)
    cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='trade_analyzer')
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe_ratio')
    cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')

    print("正在加载大盘数据...")
    market_df_copy = market_df.copy()
    market_df_copy.set_index('date', inplace=True)
    market_data_feed = bt.feeds.PandasData(
        dataname=market_df_copy,
        fromdate=start_date_for_data_obj,
        todate=end_date_obj,
        name='market_index'
    )
    cerebro.adddata(market_data_feed, name='market_index')

    print(f"正在加载 {len(hot_stocks)} 只股票数据...")
    loaded_stocks = 0
    for symbol, stock_df in tqdm(hot_stocks_dfs.items(), desc="加载股票数据"):
        if stock_df.empty or len(stock_df) < 125: 
            continue
        try:
            stock_df.set_index('date', inplace=True)
            data_feed = bt.feeds.PandasData(
                dataname=stock_df,
                fromdate=start_date_for_data_obj,
                todate=end_date_obj,
                name=symbol
            )
            cerebro.adddata(data_feed, name=symbol)
            loaded_stocks += 1
        except Exception as e:
            print(f"加载 {symbol} 数据时出错: {e}")
            continue

    print(f"成功加载 {loaded_stocks} 只股票数据进入回测引擎。")
    if loaded_stocks == 0:
        print("\n--- 未能加载任何股票数据到回测引擎，跳过此次回测。 ---")
        return create_empty_report(
            momentum_threshold, 
            max_pullback_pct, 
            trailing_stop_pct, 
            trade_size, 
            bt_settings['initial_cash'],
            'No stocks loaded'
        ), 0

    results = cerebro.run()
    
    # --- 结果分析 ---
    final_value = cerebro.broker.getvalue()
    initial_value = bt_settings['initial_cash']
    total_return = (final_value - initial_value) / initial_value * 100
    
    analyzer_results = results[0].analyzers
    trade_analysis = analyzer_results.trade_analyzer.get_analysis()
    
    total_trades = trade_analysis.total.total if trade_analysis.total else 0
    win_count = trade_analysis.won.total if trade_analysis.won else 0
    loss_count = trade_analysis.lost.total if trade_analysis.lost else 0
    
    win_rate = (win_count / total_trades * 100) if total_trades > 0 else 0
    
    total_profit = trade_analysis.won.pnl.total if trade_analysis.won.pnl else 0
    total_loss = trade_analysis.lost.pnl.total if trade_analysis.lost.pnl else 0
    
    avg_win = total_profit / win_count if win_count > 0 else 0
    avg_loss = abs(total_loss / loss_count) if loss_count > 0 else 0
    
    profit_loss_ratio = avg_win / avg_loss if avg_loss > 0 else float('inf')
    
    sharpe_ratio = analyzer_results.sharpe_ratio.get_analysis()['sharperatio'] or 0
    
    report = {
        'momentum_threshold': momentum_threshold,
        'max_pullback_pct': max_pullback_pct,
        'trailing_stop_pct': trailing_stop_pct,
        'trade_size': trade_size,
        'Total Return (%)': round(total_return, 2),
        'Sharpe Ratio': round(sharpe_ratio, 3),
        'Win Rate (%)': round(win_rate, 2),
        'Profit/Loss Ratio': round(profit_loss_ratio, 2),
        'Total Trades': total_trades,
        'Final Portfolio Value': round(final_value, 2),
        'Note': 'Success'
    }
    
    # 在单次回测模式下，保存详细的交易日志
    if output_settings.get('save_trade_log', False) and not config.get('optimization_params', {}).get('enabled', False):
        if hasattr(results[0], 'trade_log') and results[0].trade_log:
            trade_df = pd.DataFrame(results[0].trade_log)
            # 使用新的通用保存函数
            save_report(trade_df, output_settings['trade_log_path'])
            # 打印消息由save_report函数内部处理
        else:
            print("未生成任何交易，故不保存交易日志。")
    
    return report, len(hot_stocks) if 'hot_stocks' in locals() else 0


def generate_signals_for_webapp(config, force_recache=False):
    """
    一个专门为WebApp设计的函数，用于生成并返回买入信号。
    它封装了数据准备和信号生成的整个流程。
    该函数会自己处理BaoStock的登录和登出。
    """
    try:
        # 1. 数据准备 (该函数内部会调用bs.login())
        print("--- [Engine] WebApp请求：开始准备数据... ---")
        stock_codes, market_df_full = prepare_data_and_get_codes(config, force_recache=force_recache)
        if not stock_codes or market_df_full is None or market_df_full.empty:
            raise ValueError("引擎未能成功准备数据")
        print("--- [Engine] WebApp请求：数据准备完毕。 ---")

        # 2. 调用信号生成
        print("--- [Engine] WebApp请求：开始生成信号... ---")
        strategy_params = config['default_strategy_params']
        bt_settings = config['backtest_settings']
        
        _MAIN_DIR = os.path.dirname(os.path.abspath(__file__))
        CACHE_DIR = os.path.join(_MAIN_DIR, 'feather_cache')

        # 日期转换
        start_date_obj = pd.to_datetime(bt_settings['start_date'])
        # PREHEAT_TRADING_DAYS 应该与 run_backtest 中保持一致
        PREHEAT_TRADING_DAYS = 80 
        start_date_for_data_obj = start_date_obj - pd.tseries.offsets.BDay(PREHEAT_TRADING_DAYS)

        # 修复Bug：只挑选信号生成函数需要的参数，而不是传递整个字典
        # 这可以避免传入'sell_on_long_upper_shadow_pct'等多余参数
        params_for_signal_gen = {
            'momentum_threshold': strategy_params['momentum_threshold'],
            'max_pullback_pct': strategy_params['max_pullback_pct'],
            'ma_divergence_20_max': strategy_params['ma_divergence_20_max'],
            'ma_divergence_60_max': strategy_params['ma_divergence_60_max'],
            'rsi_overbought_threshold': strategy_params['rsi_overbought_threshold'],
            'precise_pullback_max_dist_pct': strategy_params['precise_pullback_max_dist_pct'],
            'market_rsi_max': strategy_params['market_rsi_max']
        }

        buy_signals, _, _ = signal_generator.generate_buy_signals_with_pandas(
            all_stock_codes=stock_codes,
            market_df=market_df_full.copy(),
            cache_dir=CACHE_DIR,
            start_date_str=bt_settings['start_date'],
            end_date_str=bt_settings['end_date'],
            max_price=bt_settings.get('max_price', 9999),
            required_start_date_obj=start_date_for_data_obj.date(),
            **params_for_signal_gen
        )
        print(f"--- [Engine] WebApp请求：信号生成完毕。 ---")
        return buy_signals

    finally:
        # 确保在操作结束后登出
        print("--- [Engine] WebApp请求处理完毕，登出BaoStock。 ---")
        bs.logout()


def _get_param_combinations(config):
    """
    根据配置生成参数组合。
    如果优化关闭，则返回包含默认参数的单元素列表。
    如果优化开启，则返回所有参数的笛卡尔积组合。
    """
    opt_params_config = config.get('optimization_params', {})
    if not opt_params_config.get('enabled', False):
        # 优化关闭，返回默认参数
        default_params = config.get('default_strategy_params', {})
        return [default_params]

    # 优化开启，生成参数组合
    param_grid = opt_params_config.get('params', {})
    
    # 从param_grid创建迭代器
    keys, values = zip(*param_grid.items())
    combinations = [dict(zip(keys, v)) for v in itertools.product(*values)]
    
    print(f"\n--- 参数优化已开启，将执行 {len(combinations)} 次回测 ---")
    return combinations


def prepare_data_and_get_codes(config, force_recache=False):
    """
    封装所有数据准备和获取逻辑的函数。
    
    :param config: 全局配置字典
    :return: (stock_codes, market_df_full)
    """
    # --- 修复：确保CACHE_DIR路径是绝对的，不依赖于CWD ---
    # 获取 main.py 文件所在的目录
    _MAIN_DIR = os.path.dirname(os.path.abspath(__file__))
    # 构造到 feather_cache 的绝对路径
    CACHE_DIR = os.path.join(_MAIN_DIR, 'feather_cache')

    bt_settings = config['backtest_settings']
    data_settings = config['data_settings']

    # 登录 BaoStock
    bs.login()

    # 如果配置要求，则强制清理缓存
    if force_recache or data_settings.get('force_recache', False):
        print("\n--- 检测到需要清理缓存 ---")
        if os.path.exists(CACHE_DIR):
            try:
                shutil.rmtree(CACHE_DIR)
                print(f"缓存目录 '{CACHE_DIR}' 已被成功删除。")
            except OSError as e:
                print(f"删除缓存时出错: {e}")
                sys.exit(1)
        os.makedirs(CACHE_DIR, exist_ok=True)
        print(f"已重新创建空缓存目录 '{CACHE_DIR}'。")

    # 获取成分股代码列表
    print("\n--- 阶段一：获取成分股列表 ---")
    stock_codes = data_provider.get_stock_list(source=data_settings.get('index_constituent_symbol', 'zz500'))
    if not stock_codes:
        print("错误：未能从BaoStock获取成分股列表。程序终止。")
        bs.logout()
        sys.exit(1)
    print(f"成功获取 {len(stock_codes)} 只成分股。")

    # --- 修复：提前计算日期，以便用于过滤 ---
    start_date_obj = pd.to_datetime(bt_settings['start_date'])
    end_date_obj = pd.to_datetime(bt_settings['end_date'])
    PREHEAT_TRADING_DAYS = 80
    start_date_for_data_obj = start_date_obj - BDay(PREHEAT_TRADING_DAYS)

    # 在下载前进行次新股过滤
    stock_codes = data_provider.filter_sub_new_stocks(stock_codes, start_date_for_data_obj.date())
    
    # 获取并缓存大盘指数数据
    print("\n--- 阶段二：准备大盘指数数据 ---")
    market_df_full = get_market_data(
        cache_dir=CACHE_DIR,
        symbol=data_settings['market_index_symbol']
    )

    # 准备并缓存所有成分股的日线数据
    print(f"\n--- 阶段三：并行准备 {len(stock_codes)} 只成分股的日线数据 ---")
    # 优化：使用多进程并为每个子进程初始化BaoStock登录
    with ProcessPoolExecutor(max_workers=4, initializer=init_baostock_login) as executor:
        futures = {
            executor.submit(
                prepare_stock_data,
                stock_code=code,
                start_date=start_date_for_data_obj.strftime('%Y-%m-%d'),
                end_date=end_date_obj.strftime('%Y-%m-%d'),
                cache_dir=CACHE_DIR,
                adjust_flag=data_settings.get('adjust_flag', "2")
            ): code for code in stock_codes
        }
        
        for future in tqdm(as_completed(futures), total=len(stock_codes), desc="并行下载股票数据"):
            future.result() # 获取结果以触发任何潜在的异常

    return stock_codes, market_df_full


def main():
    """
    主执行函数
    """
    # 加载配置文件
    with open('config.json', 'r', encoding='utf-8') as f:
        config = json.load(f)

    # --- 阶段一：准备数据 ---
    all_stock_codes, market_df_full = prepare_data_and_get_codes(config)

    # 如果未能加载任何股票代码，则提前退出
    if not all_stock_codes:
        print("未能获取任何股票代码，程序终止。")
        sys.exit(1)

    # --- 阶段二：生成参数组合 ---
    param_combinations = _get_param_combinations(config)
    
    # --- 阶段三：执行回测 ---
    all_results = []
    total_runs = len(param_combinations)
    
    # --- 根据是否开启多进程优化选择不同执行方式 ---
    opt_params_config = config.get('optimization_params', {})
    use_multiprocessing = opt_params_config.get('enabled', False) and opt_params_config.get('multiprocessing', False)

    if use_multiprocessing:
        print(f"\n--- 启动多进程参数优化，共 {total_runs} 组参数 ---")
        with ProcessPoolExecutor(initializer=init_baostock_login) as executor:
            futures = {executor.submit(run_backtest, config, all_stock_codes, market_df_full.copy(), **params): params for params in param_combinations}
            
            for future in tqdm(as_completed(futures), total=total_runs, desc="优化进度"):
                params = futures[future]
                try:
                    report, num_hot_stocks = future.result()
                    report['hot_stocks_count'] = num_hot_stocks
                    # 将参数也添加到报告中
                    report.update(params)
                    all_results.append(report)
                except Exception as e:
                    print(f"参数组 {params} 运行出错: {e}")
    else:
        # --- 单次回测或串行优化模式 ---
        for i, params in enumerate(param_combinations):
            print(f"\n--- 正在执行第 {i+1}/{total_runs} 次回测 ---")
            
            report, num_hot_stocks = run_backtest(
                config=config,
                stock_codes=all_stock_codes,
                market_df=market_df_full.copy(),
                **params
            )
            report['hot_stocks_count'] = num_hot_stocks
            # 将参数也添加到报告中
            report.update(params)
            all_results.append(report)

    # --- 阶段四：保存并打印结果 ---
    output_settings = config['output_settings']
    if all_results:
        results_df = pd.DataFrame(all_results)
        # 根据总回报率排序
        # 在多参数运行时，确保排序的列存在
        if 'Total Return (%)' in results_df.columns:
            results_df.sort_values(by='Total Return (%)', ascending=False, inplace=True)
        
        # 确定保存路径
        report_path = opt_params_config.get('results_path', output_settings['report_path']) if use_multiprocessing else output_settings['report_path']
        
        save_report(results_df, report_path)
        
        # 打印最佳结果
        if not results_df.empty:
            print("\n--- 最佳回测结果 ---")
            # 确保按Total Return排序后再取第一个
            best_result = results_df.sort_values(by='Total Return (%)', ascending=False).iloc[0].to_dict()
            for key, value in best_result.items():
                print(f"{key}: {value}")
    else:
        print("没有执行任何回测或未产生任何结果。")

    # 退出 baostock
    bs.logout()
    print("\n所有回测完成。")


if __name__ == '__main__':
    main() 